import torch
import clip
import os
from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True


class ImageClassifier:
    def __init__(self, device=None, text_batch_size: int = 1024, cache_path: str = "labels.pt"):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.text_batch_size = text_batch_size
        self.cache_path = cache_path

        # 模型保持 float32
        self.model, self.preprocess = clip.load("ViT-B/32", self.device)
        self.model.eval()
        print(f"[CLIP] loaded ViT-B/32 on {self.device} (model float32)")

        self.text_features = None
        self.labels = None

    def build_label_features(self, labels: list[str], force_rebuild: bool = False):
        """
        计算并缓存所有 label 的特征
        """
        if os.path.exists(self.cache_path) and not force_rebuild:
            print(f"[CLIP] loading cached label features from {self.cache_path}")
            data = torch.load(self.cache_path)
            self.labels = data["labels"]
            self.text_features = data["features"].to(self.device)
            return

        print(f"[CLIP] building label features for {len(labels)} labels ...")
        all_features = []
        with torch.no_grad():
            for i in range(0, len(labels), self.text_batch_size):
                print(f"[CLIP] processing batch {i}/{len(labels)} -- {int(i / self.text_batch_size)}")
                batch_labels = labels[i : i + self.text_batch_size]
                text_tokens = clip.tokenize(batch_labels).to(self.device)
                batch_features = self.model.encode_text(text_tokens)
                batch_features /= batch_features.norm(dim=-1, keepdim=True)
                all_features.append(batch_features)

        self.labels = labels
        self.text_features = torch.cat(all_features, dim=0)

        # 保存到文件
        torch.save({"labels": labels, "features": self.text_features.cpu()}, self.cache_path)
        print(f"[CLIP] saved label features to {self.cache_path}")

    def classify(self, image_path: str) -> tuple[str, float]:
        """
        仅分类（必须先 build_label_features）
        """
        if self.text_features is None:
            raise RuntimeError("You must call build_label_features(labels) first!")

        with torch.no_grad():
            image = self.preprocess(Image.open(image_path)).unsqueeze(0).to(self.device)
            image = image.half()
            image_features = self.model.encode_image(image.float())
            image_features /= image_features.norm(dim=-1, keepdim=True)

            sims = (image_features @ self.text_features.T).squeeze(0)
            score, idx = torch.max(sims, dim=0)

            return self.labels[idx.item()], score.item()
